import numpy as np

class Strategy:
    """
    策略模块 (Strategy Module) - V4 (修改为初始比率 × 线性调整因子)
    
    核心职责:
    1. 接收每日数据（包含指数点位、股债利差、ETF收盘价）。
    2. 基于一个双因子（股债利差偏离度、ETF价格）的线性模型生成调整因子。
    3. 将调整因子与初始交易比率相乘，得到最终的交易信号。
    4. 提供一个独立的接口，用于查询基于当前指数点位的最大允许仓位。
    
    策略逻辑:
    - 交易信号生成: 当股债利差触发阈值时，根据偏离度和当时的ETF价格，
      线性计算出一个调整因子，与初始交易比率相乘得到最终交易比例。
    - 风险控制: 独立于交易信号，根据指数分段函数提供最大仓位限制。
    
    新公式:
    最终交易比率 = 初始交易比率 × [基础因子 + (偏离度 × spread_factor) + (etf_price × etf_factor)]
    """

    def __init__(self, params):
        """
        初始化策略处理器。

        :param params: 字典，包含策略所需的全部参数。
                       - buy_threshold (float): 触发买入的股债利差阈值。
                       - sell_threshold (float): 触发卖出的股债利差阈值。
                       - initial_buy_ratio (float): 初始买入比率，如0.003表示0.3%。
                       - initial_sell_ratio (float): 初始卖出比率，如0.003表示0.3%。
                       - buy_base_factor (float): 买入时的基础调整因子（新的截距项）。
                       - sell_base_factor (float): 卖出时的基础调整因子（新的截距项）。
                       - buy_spread_factor (float): 买入时股债利差偏离度的权重因子。
                       - sell_spread_factor (float): 卖出时股债利差偏离度的权重因子。
                       - buy_etf_factor (float): 买入时ETF价格的权重因子。
                       - sell_etf_factor (float): 卖出时ETF价格的权重因子。
                       - max_trade_ratio (float): 单次交易的最大比例上限。
                       - index_thresholds (list): 仓位风控的指数点位阈值。
                       - max_ratios (list): 对应的最大仓位比例。
        """
        self.params = params

    def get_max_allowed_position(self, daily_data):
        """
        根据指数点位，使用分段函数确定最大允许仓位。
        这是一个独立的风控检查方法，供回测引擎调用。
        
        :param daily_data: 包含 'close' 属性（指数收盘价）的每日数据对象。
        :return: 最大允许的持仓比例 (0.0 to 1.0)。
        """
        close_price = daily_data.close # 注意：风控依然使用指数点位
        thresholds = self.params.get('index_thresholds', [])
        ratios = self.params.get('max_ratios', [1.0])

        for i, threshold_value in enumerate(thresholds):
            if close_price < threshold_value:
                return ratios[i]
                
        return ratios[len(thresholds)]

    def generate_trade_signal(self, daily_data):
        """
        根据每日数据生成交易信号（即期望的仓位变动比例）。

        :param daily_data: 包含 'spread' 和 'etf_close' 属性的每日数据对象。
        :return: 浮点数。正数表示买入比例，负数表示卖出比例，0表示无操作。
        """
        spread = daily_data.spread
        etf_price = daily_data.etf_close
        
        buy_threshold = self.params.get('buy_threshold', 6.0)
        sell_threshold = self.params.get('sell_threshold', 3.0)

        trade_signal = 0.0
        
        if spread >= buy_threshold:
            # 计算买入偏离度
            deviation = spread - buy_threshold
            # 计算买入交易比例 = 初始比率 × 调整因子
            trade_ratio = self._calculate_adjusted_trade_ratio(deviation, etf_price, 'buy')
            trade_signal = trade_ratio
            
        elif spread <= sell_threshold:
            # 计算卖出偏离度 (确保为正数)
            deviation = sell_threshold - spread
            # 计算卖出交易比例 = 初始比率 × 调整因子
            trade_ratio = self._calculate_adjusted_trade_ratio(deviation, etf_price, 'sell')
            # 卖出信号为负数
            trade_signal = -trade_ratio
            
        return trade_signal

    def _calculate_adjusted_trade_ratio(self, deviation, etf_price, side):
        """
        根据新的公式计算交易比例：初始交易比率 × 线性调整因子
        
        调整因子 = 基础因子 + 因子1*权重1 + 因子2*权重2
        最终交易比率 = 初始交易比率 × 调整因子
        """
        # 获取初始交易比率
        if side == 'buy':
            initial_ratio = self.params.get('initial_buy_ratio', 0.003)  # 默认0.3%
            base_factor = self.params.get('buy_base_factor', 1.0)        # 默认基础因子为1
        else: # sell
            initial_ratio = self.params.get('initial_sell_ratio', 0.003)  # 默认0.3%
            base_factor = self.params.get('sell_base_factor', 1.0)        # 默认基础因子为1

        # 从参数中获取两个因子的权重
        spread_factor = self.params.get('spread_factor', 0.1)      # 调整默认值，因为现在是乘法
        etf_factor = self.params.get('etf_factor', 0.0001)        # 调整默认值，因为现在是乘法
        
        # 从参数中获取单次交易上限，默认为1%
        max_trade_ratio = self.params.get('max_trade_ratio', 0.01)

        # 【核心公式修改】计算线性调整因子
        adjustment_factor = base_factor + (deviation * spread_factor) + (etf_price * etf_factor)
        
        # 确保调整因子为正数（避免负的调整因子）
        adjustment_factor = max(0, adjustment_factor)
        
        # 计算最终交易比率 = 初始比率 × 调整因子
        calculated_ratio = initial_ratio * adjustment_factor

        # 应用交易上限
        final_ratio = min(calculated_ratio, max_trade_ratio)
                
        return final_ratio